use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};

/// 零拷贝内存池
pub struct ZeroCopyMemoryPool {
    buffer: Arc<Vec<u8>>,
    offset: AtomicUsize,
    capacity: usize,
}

impl ZeroCopyMemoryPool {
    /// 创建新的零拷贝内存池
    pub fn new(capacity: usize) -> Self {
        let buffer = vec![0u8; capacity];
        ZeroCopyMemoryPool {
            buffer: Arc::new(buffer),
            offset: AtomicUsize::new(0),
            capacity,
        }
    }

    /// 分配内存
    pub fn allocate(&self, size: usize) -> Option<*mut u8> {
        let current_offset = self.offset.load(Ordering::SeqCst);
        
        if current_offset + size <= self.capacity {
            let new_offset = current_offset + size;
            if self.offset.compare_exchange(current_offset, new_offset, Ordering::SeqCst, Ordering::SeqCst).is_ok() {
                let ptr = self.buffer.as_ptr() as *mut u8;
                return Some(unsafe { ptr.add(current_offset) });
            }
        }
        
        None
    }

    /// 获取已使用内存
    pub fn used(&self) -> usize {
        self.offset.load(Ordering::SeqCst)
    }

    /// 获取总容量
    pub fn capacity(&self) -> usize {
        self.capacity
    }

    /// 重置内存池 - 安全释放内存
    pub fn reset(&mut self) {
        // 创建新的空buffer，旧buffer将被drop释放
        let new_buffer = vec![0u8; self.capacity];
        self.buffer = Arc::new(new_buffer);
        self.offset.store(0, Ordering::SeqCst);
    }

    /// 完全释放内存池资源
    pub fn destroy(&mut self) {
        // 通过创建空vec来释放内存
        self.buffer = Arc::new(Vec::new());
        self.offset.store(0, Ordering::SeqCst);
        self.capacity = 0;
    }
}

impl Default for ZeroCopyMemoryPool {
    fn default() -> Self {
        Self::new(1024 * 1024) // 1MB default
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_zero_copy_pool_creation() {
        let pool = ZeroCopyMemoryPool::new(1024);
        assert_eq!(pool.capacity(), 1024);
        assert_eq!(pool.used(), 0);
    }

    #[test]
    fn test_allocate_and_deallocate() {
        let pool = ZeroCopyMemoryPool::new(1024);
        
        if let Some(ptr) = pool.allocate(100) {
            assert!(!ptr.is_null());
            assert_eq!(pool.used(), 100);
        }
    }

    #[test]
    fn test_allocate_failure() {
        let pool = ZeroCopyMemoryPool::new(100);
        
        let result = pool.allocate(200);
        assert!(result.is_none());
    }

    #[test]
    fn test_reset() {
        let mut pool = ZeroCopyMemoryPool::new(1024);
        
        if let Some(ptr) = pool.allocate(100) {
            assert!(!ptr.is_null());
            assert_eq!(pool.used(), 100);
            
            pool.reset();
            assert_eq!(pool.used(), 0);
        }
    }
}